import requests
import psycopg2
from datetime import datetime, timedelta
from psycopg2 import sql
from typing import Dict, List, Optional, Any

# 获取accessToken
def get_access_token(app_key: str, app_secret: str) -> str:
    """获取钉钉访问令牌"""
    url = "https://api.dingtalk.com/v1.0/oauth2/accessToken"
    payload = {
        'appKey': app_key,
        'appSecret': app_secret
    }
    try:
        response = requests.post(url, json=payload)
        response.raise_for_status()
        return response.json().get('accessToken')
    except requests.exceptions.RequestException as e:
        raise Exception(f"获取访问令牌失败: {e}")

# 生成日期范围
def generate_date_range(start_date: datetime, end_date: datetime) -> List[str]:
    """生成从start_date到end_date的日期列表，格式为YYYY-MM-DD"""
    dates = []
    current_date = start_date
    while current_date <= end_date:
        dates.append(current_date.strftime("%Y-%m-%d"))
        current_date += timedelta(days=1)
    return dates

# 获取企业考勤排班详情
def fetch_listschedule(access_token: str, work_date: str) -> List[Dict[str, Any]]:
    """获取指定日期的考勤排班信息"""
    url = "https://oapi.dingtalk.com/topapi/attendance/listschedule"
    headers = {
        'Content-Type': 'application/json'
    }
    params = {
        'access_token': access_token
    }
    body = {
        'workDate': work_date
    }
    try:
        response = requests.post(url, headers=headers, params=params, json=body)
        response.raise_for_status()
        data = response.json()
        if data.get('errcode') == 0:
            result = data.get('result', {})
            return result.get('schedules', [])
        return []
    except requests.exceptions.RequestException as e:
        print(f"获取排班信息失败: {e}")
        return []

# 获取计划考勤时间
def get_plan_check_time(schedules: List[Dict[str, Any]], userid: str, check_type: str) -> Optional[str]:
    """获取指定用户的计划考勤时间"""
    for schedule in schedules:
        if schedule.get('userid') == userid and schedule.get('check_type') == check_type:
            return schedule.get('plan_check_time')
    return None

# 获取所有排班信息
def get_all_schedules(access_token: str, work_dates: List[str]) -> List[Dict[str, Any]]:
    """获取并处理多个日期的排班信息"""
    all_schedules = []

    for work_date in work_dates:
        daily_schedules = fetch_listschedule(access_token, work_date)
        user_schedules = {}

        for schedule in daily_schedules:
            userid = schedule.get('userid')
            if not userid:
                continue

            if userid not in user_schedules:
                user_schedules[userid] = {
                    'userid': userid,
                    'class_id': schedule.get('class_id',None),
                    'check_on_duty_time': get_plan_check_time(daily_schedules, userid, 'OnDuty'),
                    'check_off_duty_time': get_plan_check_time(daily_schedules, userid, 'OffDuty'),
                    'work_date': work_date
                }

        all_schedules.extend(list(user_schedules.values()))

    return all_schedules

# 查询当前用户数据是否存在
def exists_schedules(schedules: List[Dict[str, Any]], userid: str) -> bool:
    """检查用户排班数据是否存在"""
    return any(schedule.get('userid') == userid for schedule in schedules)

# 连接到数据库并插入数据
def insert_schedule_data(schedules: List[Dict[str, Any]]) -> None:
    """将排班数据插入数据库"""
    if not schedules:
        print("没有排班数据需要插入")
        return

    conn = None
    try:
        # 从环境变量或配置文件获取数据库连接信息
        conn = psycopg2.connect(
            dbname="metabase",
            user="postgres",
            password="postgres",
            host="172.16.5.226",
            port="32685"
        )

        with conn.cursor() as cur:
            # 按日期分组处理数据
            date_groups = {}
            for schedule in schedules:
                work_date = schedule['work_date']
                if work_date not in date_groups:
                    date_groups[work_date] = []
                date_groups[work_date].append(schedule)

            # 对每个日期分别处理
            for work_date, date_schedules in date_groups.items():
                # 清空当前日期的数据
                cur.execute(
                    "DELETE FROM tp.project_member_shift_record WHERE work_date::date = %s",
                    (work_date,)
                )

                # 准备数据
                data = [
                    (
                        schedule['userid'],
                        schedule['check_on_duty_time'],
                        schedule['check_off_duty_time'],
                        f"{work_date} 00:00:00",
                        'dingding',
                        schedule['class_id']
                    )
                    for schedule in date_schedules
                ]

                # 使用execute_values进行批量插入
                from psycopg2.extras import execute_values
                insert_query = sql.SQL("""
                                       INSERT INTO tp.project_member_shift_record
                                           (userid, check_on_duty_time, check_off_duty_time, work_date,data_source,class_id)
                                       VALUES %s
                                       """)
                execute_values(cur, insert_query, data)

        conn.commit()
        print(f"成功插入 {len(schedules)} 条排班数据")

    except psycopg2.Error as e:
        print(f"数据库操作失败: {e}")
        if conn:
            conn.rollback()
    finally:
        if conn:
            conn.close()

# 主函数
def main():
    """主程序入口"""
    try:
        # 获取access_token
        app_key = 'dingtyw4aiyucyykv1rr'
        app_secret = 'fkGxuNZ653SH_CyAdT7mhSCfHy40idkCAUgc_6bSk_QceBjBfg79taOBmVnLGyPO'
        access_token = get_access_token(app_key, app_secret)

        if not access_token:
            raise Exception("获取访问令牌失败")

        # 计算日期范围：今天之前一个月
        end_date = datetime.now()
        start_date = end_date - timedelta(days=30)

        # 生成日期列表
        work_dates = generate_date_range(start_date, end_date)

        # 获取企业考勤排班详情
        schedules = get_all_schedules(access_token, work_dates)

        if not schedules:
            print("当前无排班数据")
        else:
            insert_schedule_data(schedules)
            print("排班数据已成功插入数据库")

    except Exception as e:
        print(f"程序执行出错: {e}")

if __name__ == "__main__":
    main()